//! @file   lc_avl_tree.h
//! @brief  
//! @author Sayan Chaliha
//! @date   October, 2011
//!
//! 
#ifndef _LC_AVL_TREE_H
#define _LC_AVL_TREE_H

#include "libcyan_internal.h"
#include "lc_types.h"
#include "lc_numeric_limits.h"
#include "lc_object.h"
#include "lc_pair.h"
#include "lc_exception.h"

#define LC_AVL_TREE_USE_RECURSION 1

namespace cyan {
  namespace internal {
    struct CYAN_HIDDEN Direction {
      enum ReservedDirection {
        LEFT = 0,
        RIGHT
      };

      Direction() : direction_(LEFT) {}

      Direction(const Direction& other) {
        this->direction_ = other.direction_;
      }

      Direction(ReservedDirection d) : direction_(d) {}

      template <typename T>
      Direction(T d) {
        this->set<T>(d);
      }

      inline operator Int() const {
        return this->get<Int>();
      }

      template <typename T>
      inline void
      set(typename EnableIf<IsIntegral<T>::value, T>::Type d) {
        assertGtOrEq(d, LEFT);
        assertLtOrEq(d, RIGHT);
        this->direction_ = static_cast<ReservedDirection>(d);
      }

      template <typename T>
      inline typename EnableIf<IsIntegral<T>::value, T>::Type
      get() const {
        return implicit_cast<T>(this->direction_);
      }

      inline Direction& operator =(const Direction& other) {
        this->direction_ = other.direction_;
        return *this;
      }

      template <typename T>
      inline Direction&
      operator =(typename EnableIf<IsIntegral<T>::value, T>::Type d) {
        this->set<T>(d);
        return *this;
      }

      inline Direction operator ~() const {
        return Direction(1 - this->direction_);
      }

      inline Bool operator ==(const Direction& other) const {
        return this->direction_ == other.direction_;
      }

      inline Bool operator !=(const Direction& other) const {
        return this->direction_ != other.direction_;
      }

      template <typename T>
      inline Bool operator ==(typename EnableIf<IsIntegral<T>::value, T>
          ::Type d) const {
        return this->get<T>() == d;
      }

      template <typename T>
      inline Bool operator !=(typename EnableIf<IsIntegral<T>::value, T>
          ::Type d) const {
        return this->get<T>() != d;
      }

    private:
      ReservedDirection direction_;
    };

    struct CYAN_HIDDEN TreeNodeBase {
      enum {
        BALANCE_EQUAL = -1,
        BALANCE_LEFT_HEAVY = Direction::LEFT,
        BALANCE_RIGHT_HEAVY = Direction::RIGHT
      };

      TreeNodeBase* subTrees[2];
#ifndef LC_AVL_TREE_USE_RECURSION
      Int balance:2;
#endif /* LC_AVL_TREE_USE_RECURSION */

      inline static SSize height(const TreeNodeBase* node) {
        if (!node) {
          return -1;
        } else {
          const SSize leftHeight = TreeNodeBase::height
              (node->subTrees[Direction::LEFT]);
          const SSize rightHeight = TreeNodeBase::height
              (node->subTrees[Direction::RIGHT]);
          return 1 + cyan::maximum(leftHeight, rightHeight);
        }
      }

      SSize height() const throw() {
        return TreeNodeBase::height(this);
      }

      SSize balanceFactor() const throw() {
        return this->height(this->subTrees[Direction::LEFT]) -
            this->height(this->subTrees[Direction::RIGHT]);
      }

      inline Bool isBalanced() const throw() {
#ifndef LC_AVL_TREE_USE_RECURSION
        return this->balance == BALANCE_EQUAL;
#else
        return this->balanceFactor() == 0;
#endif /* LC_AVL_TREE_USE_RECURSION */
      }

      TreeNodeBase* rotate(Direction dir) throw() {
        TreeNodeBase* oldRoot = this;
        TreeNodeBase* root = this->subTrees[~dir];

        if (!root)
          return this;

        oldRoot->subTrees[~dir] = root->subTrees[dir];
        root->subTrees[dir] = oldRoot;
#ifndef LC_AVL_TREE_USE_RECURSION
        oldRoot->balance = root->balance = BALANCE_EQUAL;
#endif /* LC_AVL_TREE_USE_RECURSION */

        return root;
      }
    };

    template <typename K, typename T>
    struct CYAN_HIDDEN TreeNode : public TreeNodeBase {
      Pair<K, T> data;
    };

    template <typename K, typename T>
    class CYAN_HIDDEN AvlTreeBase {
    public:
      AvlTreeBase() {
        this->initialize();
      }

      virtual ~AvlTreeBase() {
        this->clear();
      }

      inline void initialize() {
        this->root_ = nullptr;
      }

      void clear() {
        AvlTreeBase::clearSubTree(this->root_);
        this->initialize();
      }

    protected:
      static TreeNode<K, T>* getNode() {
        return reinterpret_cast<TreeNode<K, T>*>
            (::malloc(sizeof (TreeNode<K, T>)));
      }

      static void putNode(TreeNode<K, T>* node) {
        free(node);
      }

      static void destroyData(TreeNode<K, T>* node) {
        node->data.~Pair<K, T>();
      }

      static void constructData(TreeNode<K, T>* node,
          const Pair<K, T>& value) {
        node->data = *new (static_cast<void*>(&node->data)) Pair<K, T>(value);
      }

      static void clearSubTree(TreeNodeBase* root) {
        if (root) {
          AvlTreeBase::clearSubTree(static_cast<TreeNode<K, T>*>
              (root->subTrees[Direction::LEFT]));
          AvlTreeBase::clearSubTree(static_cast<TreeNode<K, T>*>
              (root->subTrees[Direction::RIGHT]));
          AvlTreeBase::destroyData(static_cast<TreeNode<K, T>*>(root));
          AvlTreeBase::putNode(static_cast<TreeNode<K, T>*>(root));
        }
      }

      TreeNodeBase* root_;
    };
  }

  template <typename K, typename T, typename C = less<K> >
  class CYAN_API AvlTree : protected internal::AvlTreeBase<K, T>,
      public Object {
  public:
    typedef Pair<K, T>                      ValueType;
    typedef K                               KeyType;
    typedef T                               MappedType;
    typedef C                               KeyCompare;
    typedef ValueType&                      Reference;
    typedef const ValueType&                ConstReference;
    typedef ValueType*                      Pointer;
    typedef const ValueType*                ConstPointer;
    typedef Size                            SizeType;
    typedef PtrDiff                         DifferenceType;

    typedef internal::AvlTreeBase<K, T>     Base;
    typedef internal::TreeNode<K, T>        Node;
    typedef internal::TreeNodeBase          NodeBase;

    void clear() {
      this->clear();
      assertEqual(this->root_, nullptr);
    }

    void insert(const KeyType& key, const MappedType& value) {
      AvlTree::insert(&this->root_, makePair(key, value));
      assertNotNull(this->root_);
    }

    void insert(const ValueType& value) {
      AvlTree::insert(&this->root_, value);
      assertNotNull(this->root_);
    }

    void erase(const KeyType& key) {
      AvlTree::erase(&this->root_, key);
    }

    MappedType find(const KeyType& key) const {
      return AvlTree::find(this->root_, key);
    }

    Size size() const {
      return AvlTree::nodeCount(this->root_);
    }

    Bool isEmpty() const {
      return this->root_ == nullptr;
    }

  private:
    using Base::root_;
    using Base::getNode;
    using Base::putNode;
    using Base::constructData;
    using Base::destroyData;
    typedef internal::Direction Direction;

    static Node* createNode(const ValueType& value) {
      Node* node = AvlTree::getNode();

      if (!node)
        Exception::throwException
            (OutOfMemoryException(TEXT("AvlTree::createNode")));

      try {
        AvlTree::constructData(node, value);
      } catch (...) {
        AvlTree::putNode(node);
        throw;
      }

      node->subTrees[Direction::LEFT] = nullptr;
      node->subTrees[Direction::RIGHT] = nullptr;
#ifndef LC_AVL_TREE_USE_RECURSION
      node->balance = Node::BALANCE_EQUAL;
#endif /* LC_AVL_TREE_USE_RECURSION */

      return node;
    }

    static Node* createNode() {
      Node* node = AvlTree::getNode();

      if (!node)
        Exception::throwException
            (OutOfMemoryException(TEXT("AvlTree::createNode")));

      try {
        AvlTree::constructData(node, ValueType());
      } catch (...) {
        AvlTree::putNode(node);
        throw;
      }

      node->subTrees[Direction::LEFT] = nullptr;
      node->subTrees[Direction::RIGHT] = nullptr;
#ifndef LC_AVL_TREE_USE_RECURSION
      node->balance = Node::BALANCE_EQUAL;
#endif /* LC_AVL_TREE_USE_RECURSION */

      return node;
    }

    static void rebalance(NodeBase** root) {
      if (root == nullptr)
        Exception::throwException
            (NullPointerException(TEXT("AvlTree::rebalance_")));

#ifndef LC_AVL_TREE_USE_RECURSION
#else
      if (*root == nullptr)
        return;

      Node* node = static_cast<Node*>(*root);
      SSize balanceFactor = node->balanceFactor();

      if (balanceFactor != 0) {
        Direction dir = balanceFactor < 0 ?
            Direction::RIGHT : Direction::LEFT;

        NodeBase** subTree = &node->subTrees[dir];
        Node* subTreeRoot = static_cast<Node*>(*subTree);
        SSize subTreeBalanceFactor = subTreeRoot->balanceFactor();
        Direction subDir = subTreeBalanceFactor < 0 ?
            Direction::RIGHT : Direction::LEFT;

        if (dir == subDir) {
          *root = node->rotate(~subDir);
        } else {
          *subTree = subTreeRoot->rotate(~subDir);
          *root = node->rotate(~dir);
        }
      }
#endif /* LC_AVL_TREE_USE_RECURSION */
    }

    static void insert(NodeBase** root, const ValueType& value) {
      if (root == nullptr)
        Exception::throwException
            (NullPointerException(TEXT("AvlTree::insert_")));

#ifndef LC_AVL_TREE_USE_RECURSION
      NodeBase** root = &this->root_;
      Node* node = static_cast<Node*>(*root);

      while (node && node->data.first != value.first) {
        Direction dir = KeyCompare()(node->data.first, value.first);

        root = &node->subTrees[dir];
        node = static_cast<Node*>(*root);
      }

      if (node) {
        AvlTree::destroyData(node);
        AvlTree::constructData(node, value);
      } else {
        node = AvlTree::createNode(value);
        *root = node;
      }
#else
      if (*root == nullptr) {
        *root = AvlTree::createNode(value);
      } else {
        Node* node = static_cast<Node*>(*root);
        Direction dir = KeyCompare()(node->data.first, value.first);
        AvlTree::insert(&node->subTrees[dir], value);
        AvlTree::rebalance(root);
      }
#endif /* LC_AVL_TREE_USE_RECURSION */
    }

    static ValueType minValue(Node* root) {
      if (root->subTrees[Direction::LEFT] == nullptr)
        return root->data;
      else
        return AvlTree::minValue_
            (static_cast<Node*>(root->subTrees[Direction::LEFT]));
    }

    static void erase(NodeBase** root, const KeyType& key) {
      if (root == nullptr)
        Exception::throwException
            (NullPointerException(TEXT("AvlTree::delete_")));

#ifndef LC_AVL_TREE_USE_RECURSION
      NodeBase** root = &this->root_;
      Node* node = static_cast<Node*>(*root);
      NodeBase** parent = null;

      while (node && node->data.first != key) {
        Direction dir = KeyCompare()(node->data.first, key);

        parent = root;
        root = &node->subTrees[dir];
        node = static_cast<Node*>(*root);
      }

      if (node) {
      }
#else
      if (*root == nullptr)
        return;

      Node* node = static_cast<Node*>(*root);
      if (node->data.first != key) {
        Direction dir = KeyCompare()(node->data.first, key);
        AvlTree::delete_(&node->subTrees[dir], key);
      } else {
        if (node->subTrees[Direction::LEFT] == nullptr ||
            node->subTrees[Direction::RIGHT] == nullptr) {
          Direction dir = node->subTrees[Direction::LEFT] == nullptr ?
              Direction::RIGHT : Direction::LEFT;

          *root = node->subTrees[dir];
          AvlTree::destroyData(node);
          AvlTree::putNode(node);
        } else {
          ValueType successorValue = AvlTree::minValue_
              (static_cast<Node*>(node->subTrees[Direction::RIGHT]));
          AvlTree::destroyData(node);
          AvlTree::constructData(node, successorValue);
          AvlTree::delete_(&node->subTrees[Direction::RIGHT], node->data.first);
        }
        AvlTree::rebalance(root);
      }
#endif /* LC_AVL_TREE_USE_RECURSION */
    }

    static MappedType find(const NodeBase* root, const KeyType& key) {
      if (root == nullptr)
        return MappedType();

#ifndef LC_AVL_TREE_USE_RECURSION

#else
      const Node* node = static_cast<const Node*>(root);
      if (node->data.first != key) {
        Direction dir = KeyCompare()(node->data.first, key);
        return AvlTree::find(node->subTrees[dir], key);
      } else {
        return node->data.second;
      }
#endif /* LC_AVL_TREE_USE_RECURSION */
    }

    static Size nodeCount(const NodeBase* root) {
      if (root) {
        return 1 + nodeCount(root->subTrees[Direction::LEFT]) +
            nodeCount(root->subTrees[Direction::RIGHT]);
      }
      return 0;
    }
  };
}

#endif /* _LC_AVL_TREE_H */
